import os
from os.path import join, isfile
import sys
sys.path.append("..")
import argparse
import numpy as np
import pickle
import torch
from torch import nn
import pytorch3d
import pytorch3d.io as IO
import trimesh
from smplx import smplx
from smplx.smplx.utils import Struct, to_tensor, to_np
import cv2
import imageio
from utils.txt2intrinsic import txt2intrinsic
from utils.pyt3d_wrapper import Pyt3DWrapper
from utils.avi2depth import avi2depth
from utils.time_align import time_align
from utils.process_timestamps import txt_to_paried_frameids, paired_frameids_to_txt
from utils.contact import compute_contact
from utils.VTS_object import get_obj_info, get_obj_name_correspondance
from utils.visualization import save_mesh
from utils.load_smplx_params import load_multiperson_smplx_params
from utils.object_retargeting import obj_retargeting
from utils.contact import compute_contact_and_closest_point
from smplx.smplx.lbs import batch_rodrigues, vertices2joints, blend_shapes, batch_rigid_transform
from transforms3d.axangles import mat2axangle
import open3d as o3d
from optimization.utils import local_pose_to_global_orientation
from optimization.bvh2smplx import Simple_SMPLX, create_SMPLX_model
from utils.retargeting_visualization import HOI_visualization
import matplotlib.pyplot as plt
import time
from utils.simplify_mesh import simplify_mesh
from get_joints import get_joints

HAND_VERT_IDS = {
    'lthumb':		5361,
    'lindex':		4933,
    'lmiddle':		5058,
    'lring':		5169,
    'lpinky':		5286,
    'rthumb':		8079,
    'rindex':		7669,
    'rmiddle':		7794,
    'rring':		7905,
    'rpinky':		8022,  
}

expression = [[ 0.2816, -0.0337,  0.0660],
        [ 0.2301, -0.0678,  0.0553],
        [ 0.2616, -0.0645,  0.1151],
        [ 0.1371, -0.0239,  0.0633],
        [ 0.2017, -0.0290,  0.1854]]
face = [[ 0.1945, -0.0684,  0.0446],
        [ 0.2109, -0.0870,  0.0446],
        [ 0.2272, -0.0927,  0.0524],
        [ 0.2407, -0.0922,  0.0634],
        [ 0.2494, -0.0880,  0.0736],
        [ 0.2580, -0.0868,  0.0902],
        [ 0.2628, -0.0911,  0.1035],
        [ 0.2639, -0.0915,  0.1203],
        [ 0.2607, -0.0860,  0.1387],
        [ 0.2526, -0.0672,  0.1518],
        [ 0.2591, -0.0711,  0.0781],
        [ 0.2673, -0.0608,  0.0740],
        [ 0.2749, -0.0507,  0.0699],
        [ 0.2821, -0.0415,  0.0660],
        [ 0.2595, -0.0261,  0.0627],
        [ 0.2658, -0.0246,  0.0664],
        [ 0.2705, -0.0226,  0.0710],
        [ 0.2715, -0.0241,  0.0777],
        [ 0.2710, -0.0254,  0.0850],
        [ 0.2179, -0.0648,  0.0483],
        [ 0.2290, -0.0717,  0.0519],
        [ 0.2344, -0.0708,  0.0601],
        [ 0.2369, -0.0649,  0.0676],
        [ 0.2350, -0.0622,  0.0599],
        [ 0.2297, -0.0618,  0.0520],
        [ 0.2541, -0.0639,  0.1022],
        [ 0.2593, -0.0703,  0.1095],
        [ 0.2627, -0.0702,  0.1187],
        [ 0.2597, -0.0634,  0.1294],
        [ 0.2636, -0.0610,  0.1176],
        [ 0.2599, -0.0615,  0.1086],
        [ 0.2475,  0.0008,  0.0558],
        [ 0.2589, -0.0035,  0.0574],
        [ 0.2672, -0.0065,  0.0642],
        [ 0.2707, -0.0057,  0.0698],
        [ 0.2731, -0.0065,  0.0757],
        [ 0.2734, -0.0033,  0.0866],
        [ 0.2697,  0.0024,  0.0978],
        [ 0.2732,  0.0046,  0.0890],
        [ 0.2737,  0.0046,  0.0773],
        [ 0.2718,  0.0047,  0.0709],
        [ 0.2673,  0.0043,  0.0651],
        [ 0.2576,  0.0031,  0.0588],
        [ 0.2474,  0.0007,  0.0566],
        [ 0.2637,  0.0006,  0.0659],
        [ 0.2668,  0.0010,  0.0714],
        [ 0.2694,  0.0008,  0.0773],
        [ 0.2694,  0.0023,  0.0975],
        [ 0.2706, -0.0019,  0.0791],
        [ 0.2684, -0.0021,  0.0725],
        [ 0.2644, -0.0023,  0.0665]]
device = "cuda:0"
smplx_model = create_SMPLX_model()
smplx_model = smplx_model.to(device)
betas = torch.zeros((1, 10)).to(device)
expression = torch.zeros((1, 10)).to(device)
global_orient = torch.zeros((1, 3)).to(device)
transl = torch.zeros((1, 3)).to(device)
body_pose = torch.randn((1, 21, 3)).to(device)
left_hand_pose = torch.zeros((1, 12)).to(device)
right_hand_pose = torch.zeros((1, 12)).to(device)
jaw_pose = torch.zeros((1, 3)).to(device)
leye_pose = torch.zeros((1, 3)).to(device)
reye_pose = torch.zeros((1, 3)).to(device)

model_path = "/share/human_model/models/smplx/SMPLX_NEUTRAL.npz"
model_data = np.load(model_path, allow_pickle=True)
data_struct = Struct(**model_data)
shapedirs = data_struct.shapedirs
v_template = data_struct.v_template
J_regressor = data_struct.J_regressor
parents = data_struct.kintree_table[0]
left_hand_components = data_struct.hands_componentsl[:12]
right_hand_components = data_struct.hands_componentsr[:12]

model = smplx_model(betas=betas, expression=expression, global_orient=global_orient, transl=transl, body_pose=body_pose, left_hand_pose=left_hand_pose, right_hand_pose=right_hand_pose, jaw_pose=jaw_pose, leye_pose=leye_pose, reye_pose=reye_pose)
joints_my = get_joints(global_orient, betas, body_pose, transl, left_hand_pose, right_hand_pose, left_hand_components, right_hand_components, shapedirs, v_template, J_regressor, parents)[0]

vert = model.vertices[0].detach().cpu().numpy()
face = model.faces.detach().cpu().numpy()

joints = model.joints[0]
print(joints.shape)
index = [10, 11, 39, 27, 30, 36, 33, 54, 42, 45, 51, 48]
index_l = []
# print(joints_my)
# print(joints[0:55])
# joint_1 = joints_my[0:55]
joint_1 = joints_my[61:]
joint_2 = joints[66:76]
# joint_2 = joints[55:71]

print(joints[55:60], joints[76:127])
print(vert[list(HAND_VERT_IDS.values())])

print(joint_1.shape, joint_2.shape)
print(torch.where(joint_1 != joint_2))
print(joint_1, joint_2)

# idx = [15, 22, 23, 24]
# print(joint_1[idx], joint_2[idx])
# print(transl)

# print(joints[index, :])
# print(joints[60:76, :])
# print(joints[39, ], joints[38, ])


# mesh = trimesh.Trimesh(vertices=vert, faces=face)
# mesh_txt = trimesh.exchange.obj.export_obj(mesh, include_normals=False, include_color=False, include_texture=False, return_texture=False, write_texture=False, resolver=None, digits=8)
# with open("test.obj", "w") as fp:
#     fp.write(mesh_txt)